import torch

def merge_channels(input_channels, output_channels, image):

    if input_channels % output_channels != 0:
        print("Error: 整除不了")
        return None


    channels_per_group = input_channels // output_channels
    import torch

    def merge_channels(input_channels, output_channels, image):
        # 检查输入通道数能否被输出通道数整除
        if input_channels % output_channels != 0:
            print("Error: 输入通道数不能被输出通道数整除")
            return None

        # 计算每个输出通道需要合并的输入通道数
        channels_per_group = input_channels // output_channels

        # 将输入图像沿最后一个维度分割成输入通道数个通道
        image_channels = torch.split(image, 1, dim=-1)

        # 创建空列表存储合并后的通道
        merged_channels = []

        # 遍历输出通道数
        for i in range(output_channels):
            # 创建一个全零张量
            merged_channel = torch.zeros_like(image_channels[0])
            # 对于每个输出通道，将相邻的输入通道相加
            for j in range(channels_per_group):
                index = i * channels_per_group + j
                merged_channel += image_channels[index]
            # 将合并后的通道添加到列表中
            merged_channels.append(merged_channel)

        # 将合并后的通道沿最后一个维度连接
        merged_image = torch.cat(merged_channels, dim=-1)
        return merged_image

    # 输入通道数和输出通道数
    input_channels = 6
    output_channels = 2

    # 创建随机张量作为输入图像
    image = torch.rand(255, 255, input_channels)

    # 打印输入通道数
    print("输入通道数:", image.shape)

    # 使用 merge_channels 函数合并通道
    result = merge_channels(input_channels, output_channels, image)

    # 打印输出通道数
    print("输出通道数:", result.shape)

    image_channels = torch.split(image, 1, dim=-1)


    merged_channels = []


    for i in range(output_channels):

        merged_channel = torch.zeros_like(image_channels[0])

        for j in range(channels_per_group):
            index = i * channels_per_group + j
            merged_channel += image_channels[index]

        merged_channels.append(merged_channel)


    merged_image = torch.cat(merged_channels, dim=-1)
    return merged_image


input_channels = 100
output_channels = 2


image = torch.rand(255, 255, input_channels)

print("输入通道数:", image.shape)
result = merge_channels(input_channels, output_channels, image)
print("输出通道数:", result.shape)
